# 导入必要的库
import cv2  # OpenCV库，用于图像处理
import matplotlib.pyplot as plt  # Matplotlib库，用于绘制图像
import numpy as np  # NumPy库，用于数值计算
from segment_anything import sam_model_registry, SamPredictor  # Segment Anything模型的相关模块

# 定义一个函数来显示分割掩码
def show_mask(mask, ax, random_color=False):
    """
    显示分割掩码。
    参数:
    - mask: 分割掩码（numpy数组）- ax: matplotlib的Axes对象，用于绘图 - random_color: 是否使用随机颜色，默认为False
    """
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)  # 随机生成颜色
    else:
        color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])  # 使用固定的蓝色半透明颜色
    
    h, w = mask.shape[-2:]  # 获取掩码的高度和宽度
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)  # 将掩码应用到颜色上
    ax.imshow(mask_image)  # 在指定的Axes对象上显示掩码图像

# 定义一个函数来显示点（正点和负点）
def show_points(coords, labels, ax, marker_size=375):
    """
    显示点（正点和负点）。
    参数:
    - coords: 坐标数组（numpy数组）- labels: 标签数组（1表示正点，0表示负点）- ax: matplotlib的Axes对象，用于绘图
    - marker_size: 点的大小，默认为375
    """
    pos_points = coords[labels == 1]  # 获取正点坐标
    neg_points = coords[labels == 0]  # 获取负点坐标
    # 绘制正点
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
    # 绘制负点
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)

# 设置SAM模型的检查点路径、设备类型和模型类型
sam_checkpoint = r'F:\homework\ai\Model Checkpoints\sam_vit_h_4b8939.pth'  # 模型权重文件路径
device = "cuda"  # 使用GPU进行加速
model_type = "default"  # 使用默认的模型类型
# 加载并初始化SAM模型
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)  # 根据模型类型和检查点加载模型
sam.to(device=device)  # 将模型移动到指定设备（CPU或GPU）

# 创建预测器实例
predictor = SamPredictor(sam)

# 读取输入图像并转换颜色空间
image = cv2.imread(r"F:\homework\ai\demo\dog.jpg")  # 使用OpenCV读取图像
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 将BGR格式转换为RGB格式

# 设置预测器使用的图像
predictor.set_image(image)

# 定义输入点及其标签（1表示正点，0表示负点）
input_point = np.array([[1250, 800]])  # 输入点的坐标
input_label = np.array([1])  # 输入点的标签（1表示这是个正点）

# 显示原始图像以及输入点
plt.figure(figsize=(10,10))  # 创建一个新的图形窗口
plt.imshow(image)  # 显示原始图像
show_points(input_point, input_label, plt.gca())  # 在图像上显示输入点
plt.axis('on')  # 打开坐标轴
plt.show()  # 显示图像

# 使用输入点进行预测
masks, scores, logits = predictor.predict(
    point_coords=input_point,  # 输入点的坐标
    point_labels=input_label,  # 输入点的标签
    multimask_output=True,  # 允许输出多个掩码
)

# 遍历每个分割结果，并显示它们
for i, (mask, score) in enumerate(zip(masks, scores)):
    plt.figure(figsize=(10,10))  # 创建一个新的图形窗口
    plt.imshow(image)  # 显示原始图像
    show_mask(mask, plt.gca())  # 显示分割掩码
    show_points(input_point, input_label, plt.gca())  # 显示输入点
    plt.title(f"Mask {i+1}, Score: {score:.3f}", fontsize=18)  # 设置图像标题
    plt.axis('off')  # 关闭坐标轴
    plt.show()  # 显示图像